import os

import numpy as np
import pybullet as p

import constants as c
from jumping import JUMPING
from locomotion import LOCOMOTION
from motor import MOTOR
from pyrosim import pyrosim
from rotation import ROTATION
from sensor import ENV_SENSOR, LINK_SENSOR


class ROBOT:
    
    
    def __init__(self, task, target):
        
        self.__create_body()
        self.tasks = self.__initialize_task(task)
        self.__body = None
        
        
    def initialize_simulation(self):
        self.__body = self.__load_body()
        self.__link_sensors, self.__env_sensors = self.__add_sensor()
        self.__motors = self.__add_motor()
        
        
    def __create_body(self):
        
        """生成机器人模型文件"""
        
        if os.path.exists("models/body.urdf"):
            return
        
        pyrosim.Start_URDF("models/body.urdf")
        
        pyrosim.Send_Cube(name="Torso", pos=[0, 0, 1] , size=[1, 1, 1])
        
        pyrosim.Send_Joint(name="Torso_FrontLeg", parent="Torso", child="FrontLeg", type="revolute", position = "0 0.5 1", jointAxis="1 0 0")
        pyrosim.Send_Cube(name="FrontLeg", pos=[0, 0.5, 0] , size=[0.2, 1, 0.2])
        
        pyrosim.Send_Joint(name="FrontLeg_FrontLowerLeg", parent="FrontLeg", child="FrontLowerLeg", type="revolute", position = "0 1 0", jointAxis="1 0 0")
        pyrosim.Send_Cube(name="FrontLowerLeg", pos=[0, 0, -0.5] , size=[0.2, 0.2, 1])
        
        pyrosim.Send_Joint(name="Torso_BackLeg", parent="Torso", child="BackLeg", type="revolute", position = "0 -0.5 1", jointAxis="1 0 0")
        pyrosim.Send_Cube(name="BackLeg", pos=[0, -0.5, 0] , size=[0.2, 1, 0.2])
        
        pyrosim.Send_Joint(name="BackLeg_BackLowerLeg", parent="BackLeg", child="BackLowerLeg", type="revolute", position = "0 -1 0", jointAxis="1 0 0")
        pyrosim.Send_Cube(name="BackLowerLeg", pos=[0, 0, -0.5] , size=[0.2, 0.2, 1])
        
        pyrosim.Send_Joint(name="Torso_LeftLeg", parent="Torso", child="LeftLeg", type="revolute", position = "-0.5 0 1", jointAxis="0 1 0")
        pyrosim.Send_Cube(name="LeftLeg", pos=[-0.5, 0, 0] , size=[1, 0.2, 0.2])
        
        pyrosim.Send_Joint(name="LeftLeg_LeftLowerLeg", parent="LeftLeg", child="LeftLowerLeg", type="revolute", position = "-1 0 0", jointAxis="0 1 0")
        pyrosim.Send_Cube(name="LeftLowerLeg", pos=[0, 0, -0.5] , size=[0.2, 0.2, 1])
        
        pyrosim.Send_Joint(name="Torso_RightLeg", parent="Torso", child="RightLeg", type="revolute", position = "0.5 0 1", jointAxis="0 1 0")
        pyrosim.Send_Cube(name="RightLeg", pos=[0.5, 0, 0] , size=[1, 0.2, 0.2])
        
        pyrosim.Send_Joint(name="RightLeg_RightLowerLeg", parent="RightLeg", child="RightLowerLeg", type="revolute", position = "1 0 0", jointAxis="0 1 0")
        pyrosim.Send_Cube(name="RightLowerLeg", pos=[0, 0, -0.5] , size=[0.2, 0.2, 1])
        
        pyrosim.End()
        
        
    def __add_sensor(self):
        
        """给机器人添加传感器"""
        
        link_sensors = {}
        env_sensors = {}
        
        for linkName in pyrosim.linkNamesToIndices:
            link_sensors[linkName] = LINK_SENSOR(linkName)
        
        env_sensors["position"] = ENV_SENSOR("position", self.__body)
        env_sensors["orientation"] = ENV_SENSOR("orientation", self.__body)
        
        return link_sensors, env_sensors
        
        
    def __add_motor(self):
        
        """给机器人添加动力组件"""
        
        motors = {}
        for jointName in pyrosim.jointNamesToIndices:
            motors[jointName] = MOTOR(jointName)
        return motors
        
        
    def __load_body(self):
        
        """加载机器人模型文件"""
        quaternionAngle = p.getQuaternionFromEuler([0, 0, 0])
        body = p.loadURDF("models/body.urdf", baseOrientation=quaternionAngle)
        pyrosim.Prepare_To_Simulate("models/body.urdf")
        return body
        
         
    def __initialize_task(self, task):
        
        tasks = {}
        tasks["rotation"] = None
        tasks["locomotion"] = None
        tasks["jumping"] = None
        if task == "locomotion":
            tasks["locomotion"] = LOCOMOTION()
        
        elif task == "rotation":
            tasks["rotation"] = ROTATION()
        
        elif task == "jumping":
            tasks["jumping"] = JUMPING()
        
        else:
            tasks["rotation"] = ROTATION()
            tasks["locomotion"] = LOCOMOTION()
            tasks["jumping"] = JUMPING()
            
        return tasks
        
        
    def __sense(self):
        
        link_sensor_values = np.zeros(c.link_sensor_num)
        for i, sensor in enumerate(self.__link_sensors.values()):
            link_sensor_values[i] = sensor.get_value()
            
        position = self.__env_sensors["position"].get_value()
        orientation = self.__env_sensors["orientation"].get_value()
        
        return link_sensor_values, position, orientation
    
    
    def __act(self, motor_values):
        
        for motor, motor_value in zip(self.__motors.values(), motor_values):
            motor.set_value(self.__body, motor_value)
       
    @property    
    def id_(self):
        return self.__body
        
                
    def run(self, task):
            
        if type(task) is LOCOMOTION:
            angle_limit = c.angle_limit_of_locomotion
        elif type(task) is ROTATION:
            angle_limit = c.angle_limit_of_rotation
        else:
            angle_limit = c.angle_limit_of_jumping
            
        link_sensor_values, _ ,_ = self.__sense()
        
        # if task.is_finished(position, orientation, self.__target):
        #     self.tasks.pop(0)
        #     return
            
        motor_values = task.act(link_sensor_values)
        self.__act(motor_values * angle_limit)
        
        
    def get_info(self, type):
        
        return self.__env_sensors[type].get_value()
    
    
    def save_networks(self, index):
        
        for task in self.tasks.values():
            if task is not None:
                task.save_network(index)
                
    
    def save_motor_values(self):
        
        for motor in self.__motors.values():
            for task in self.tasks.keys():
                if self.tasks[task]:
                    motor.save_values(task)
        
        
if __name__ == "__main__":
    
    robot = ROBOT()
    print(robot.__network)
